"""
评估指标模块

提供医学图像分割的各种评估指标计算
"""

import numpy as np
import torch
from typing import Union, List, Dict, Any
from scipy.spatial.distance import directed_hausdorff
from skimage.metrics import structural_similarity as ssim
import cv2


def dice_coefficient(pred: Union[np.ndarray, torch.Tensor], 
                    target: Union[np.ndarray, torch.Tensor],
                    smooth: float = 1e-6) -> float:
    """
    计算Dice系数
    
    Args:
        pred: 预测掩码 (H, W) 或 (B, H, W)
        target: 真实掩码 (H, W) 或 (B, H, W)
        smooth: 平滑参数，避免除零
    
    Returns:
        Dice系数
    """
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()
    
    # 二值化
    pred = (pred > 0.5).astype(np.float32)
    target = (target > 0.5).astype(np.float32)
    
    # 计算交集和并集
    intersection = np.sum(pred * target)
    union = np.sum(pred) + np.sum(target)
    
    dice = (2.0 * intersection + smooth) / (union + smooth)
    return float(dice)


def jaccard_index(pred: Union[np.ndarray, torch.Tensor], 
                 target: Union[np.ndarray, torch.Tensor],
                 smooth: float = 1e-6) -> float:
    """
    计算Jaccard指数（IoU）
    
    Args:
        pred: 预测掩码 (H, W) 或 (B, H, W)
        target: 真实掩码 (H, W) 或 (B, H, W)
        smooth: 平滑参数
    
    Returns:
        Jaccard指数
    """
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()
    
    # 二值化
    pred = (pred > 0.5).astype(np.float32)
    target = (target > 0.5).astype(np.float32)
    
    # 计算交集和并集
    intersection = np.sum(pred * target)
    union = np.sum(pred) + np.sum(target) - intersection
    
    jaccard = (intersection + smooth) / (union + smooth)
    return float(jaccard)


def hausdorff_distance(pred: Union[np.ndarray, torch.Tensor], 
                      target: Union[np.ndarray, torch.Tensor],
                      percentile: float = 95.0) -> float:
    """
    计算Hausdorff距离
    
    Args:
        pred: 预测掩码 (H, W)
        target: 真实掩码 (H, W)
        percentile: 百分位数，用于计算HD95
    
    Returns:
        Hausdorff距离
    """
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()
    
    # 二值化
    pred = (pred > 0.5).astype(np.uint8)
    target = (target > 0.5).astype(np.uint8)
    
    # 获取边界点
    pred_contours, _ = cv2.findContours(pred, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    target_contours, _ = cv2.findContours(target, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if len(pred_contours) == 0 or len(target_contours) == 0:
        return float('inf')
    
    # 合并所有轮廓点
    pred_points = np.vstack([contour.reshape(-1, 2) for contour in pred_contours])
    target_points = np.vstack([contour.reshape(-1, 2) for contour in target_contours])
    
    # 计算Hausdorff距离
    try:
        hd1 = directed_hausdorff(pred_points, target_points)[0]
        hd2 = directed_hausdorff(target_points, pred_points)[0]
        hd = max(hd1, hd2)
        
        # 计算HD95
        if percentile < 100.0:
            distances = []
            for p in pred_points:
                min_dist = np.min(np.linalg.norm(target_points - p, axis=1))
                distances.append(min_dist)
            distances = np.array(distances)
            hd95 = np.percentile(distances, percentile)
            return float(hd95)
        else:
            return float(hd)
    except:
        return float('inf')


def sensitivity(pred: Union[np.ndarray, torch.Tensor], 
               target: Union[np.ndarray, torch.Tensor],
               smooth: float = 1e-6) -> float:
    """
    计算敏感性（召回率）
    
    Args:
        pred: 预测掩码 (H, W) 或 (B, H, W)
        target: 真实掩码 (H, W) 或 (B, H, W)
        smooth: 平滑参数
    
    Returns:
        敏感性
    """
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()
    
    # 二值化
    pred = (pred > 0.5).astype(np.float32)
    target = (target > 0.5).astype(np.float32)
    
    # 计算TP和FN
    tp = np.sum(pred * target)
    fn = np.sum(target * (1 - pred))
    
    sensitivity = (tp + smooth) / (tp + fn + smooth)
    return float(sensitivity)


def specificity(pred: Union[np.ndarray, torch.Tensor], 
               target: Union[np.ndarray, torch.Tensor],
               smooth: float = 1e-6) -> float:
    """
    计算特异性
    
    Args:
        pred: 预测掩码 (H, W) 或 (B, H, W)
        target: 真实掩码 (H, W) 或 (B, H, W)
        smooth: 平滑参数
    
    Returns:
        特异性
    """
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()
    
    # 二值化
    pred = (pred > 0.5).astype(np.float32)
    target = (target > 0.5).astype(np.float32)
    
    # 计算TN和FP
    tn = np.sum((1 - pred) * (1 - target))
    fp = np.sum(pred * (1 - target))
    
    specificity = (tn + smooth) / (tn + fp + smooth)
    return float(specificity)


def precision(pred: Union[np.ndarray, torch.Tensor], 
              target: Union[np.ndarray, torch.Tensor],
              smooth: float = 1e-6) -> float:
    """
    计算精确率
    
    Args:
        pred: 预测掩码 (H, W) 或 (B, H, W)
        target: 真实掩码 (H, W) 或 (B, H, W)
        smooth: 平滑参数
    
    Returns:
        精确率
    """
    if isinstance(pred, torch.Tensor):
        pred = pred.detach().cpu().numpy()
    if isinstance(target, torch.Tensor):
        target = target.detach().cpu().numpy()
    
    # 二值化
    pred = (pred > 0.5).astype(np.float32)
    target = (target > 0.5).astype(np.float32)
    
    # 计算TP和FP
    tp = np.sum(pred * target)
    fp = np.sum(pred * (1 - target))
    
    precision = (tp + smooth) / (tp + fp + smooth)
    return float(precision)


def f1_score(pred: Union[np.ndarray, torch.Tensor], 
            target: Union[np.ndarray, torch.Tensor],
            smooth: float = 1e-6) -> float:
    """
    计算F1分数
    
    Args:
        pred: 预测掩码 (H, W) 或 (B, H, W)
        target: 真实掩码 (H, W) 或 (B, H, W)
        smooth: 平滑参数
    
    Returns:
        F1分数
    """
    prec = precision(pred, target, smooth)
    sens = sensitivity(pred, target, smooth)
    
    f1 = (2 * prec * sens) / (prec + sens + smooth)
    return float(f1)


def calculate_all_metrics(pred: Union[np.ndarray, torch.Tensor], 
                         target: Union[np.ndarray, torch.Tensor]) -> Dict[str, float]:
    """
    计算所有评估指标
    
    Args:
        pred: 预测掩码 (H, W) 或 (B, H, W)
        target: 真实掩码 (H, W) 或 (B, H, W)
    
    Returns:
        包含所有指标的字典
    """
    metrics = {
        'dice': dice_coefficient(pred, target),
        'jaccard': jaccard_index(pred, target),
        'hd95': hausdorff_distance(pred, target, percentile=95.0),
        'sensitivity': sensitivity(pred, target),
        'specificity': specificity(pred, target),
        'precision': precision(pred, target),
        'f1': f1_score(pred, target)
    }
    
    return metrics


def batch_metrics(predictions: List[Union[np.ndarray, torch.Tensor]], 
                 targets: List[Union[np.ndarray, torch.Tensor]]) -> Dict[str, List[float]]:
    """
    批量计算指标
    
    Args:
        predictions: 预测掩码列表
        targets: 真实掩码列表
    
    Returns:
        包含所有样本指标的字典
    """
    all_metrics = {
        'dice': [],
        'jaccard': [],
        'hd95': [],
        'sensitivity': [],
        'specificity': [],
        'precision': [],
        'f1': []
    }
    
    for pred, target in zip(predictions, targets):
        metrics = calculate_all_metrics(pred, target)
        for key, value in metrics.items():
            all_metrics[key].append(value)
    
    return all_metrics


def mean_std_metrics(metrics_dict: Dict[str, List[float]]) -> Dict[str, Dict[str, float]]:
    """
    计算指标的平均值和标准差
    
    Args:
        metrics_dict: 指标字典
    
    Returns:
        包含均值和标准差的字典
    """
    result = {}
    
    for metric_name, values in metrics_dict.items():
        values = np.array(values)
        result[metric_name] = {
            'mean': float(np.mean(values)),
            'std': float(np.std(values)),
            'min': float(np.min(values)),
            'max': float(np.max(values))
        }
    
    return result


def print_metrics_summary(metrics_dict: Dict[str, Dict[str, float]]) -> None:
    """
    打印指标摘要
    
    Args:
        metrics_dict: 指标字典
    """
    print("\n" + "="*60)
    print("评估指标摘要")
    print("="*60)
    
    for metric_name, stats in metrics_dict.items():
        print(f"{metric_name.upper():>12}: {stats['mean']:.4f} ± {stats['std']:.4f} "
              f"(范围: {stats['min']:.4f} - {stats['max']:.4f})")
    
    print("="*60)
